import torch
from CNN import CNN
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data
from openpyxl.styles.builtins import output
import matplotlib.pyplot as plt

acc_list_test = []
EPOCH=50

train=dataset.MNIST(root="mnist",train=True,transform=transforms.ToTensor(),download=True)
test=dataset.MNIST(root="mnist",train=False,transform=transforms.ToTensor(),download=True)
#进行分批加载
train_loader=torch.utils.data.DataLoader(dataset=train,batch_size=64,shuffle=True)#打乱数据顺序
test_loader=torch.utils.data.DataLoader(dataset=test,batch_size=64,shuffle=True)#打乱数据顺序

cnn=CNN();

#交叉熵损失
loss_func=torch.nn.CrossEntropyLoss()
optimizer=torch.optim.Adam(cnn.parameters(),lr=0.01)
for epoch in range(EPOCH):#训练50次
    for index,(images,labels) in enumerate(train_loader): #元组
        optimizer.zero_grad()#清空梯度
        outputs=cnn(images)#前向传播
        loss=loss_func(outputs,labels)
        loss.backward()#反向传播
        optimizer.step()
    total=0
    correct = 0
    for index2, (images,labels) in enumerate(test_loader):
        outputs=cnn(images)
        loss_test=loss_func(outputs,labels)
        maxValue=outputs.max(1)
        _,pred=outputs.max(1)
        print(pred)
        total+=labels.size(0)
        correct+=(pred==labels).sum().item()
    acc=correct/total
    acc_list_test.append(acc)
    print('[%d / %d]: 准确率: %.1f %% ' % (epoch + 1, EPOCH, 100 * acc))  # 求测试的准确率，正确数/总数
torch.save(cnn,"MNIST/mnist_model.pkl")
#画图
plt.plot(acc_list_test)
plt.title('epoch acc')
plt.xlabel('epoch')
plt.ylabel('acc')
plt.grid(True)
plt.plot(acc_list_test, 'b-', linewidth=2)
plt.ylim(0.8, 1.0)  # 限制y轴范围以突出变化
plt.show()